'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: dataset_schema.py
* @Time: 2025/9/25
* @All Rights Reserve By Brtc
'''
from flask_wtf import FlaskForm
from marshmallow import fields, pre_dump, Schema
from wtforms import StringField, IntegerField, FloatField
from wtforms.validators import DataRequired, Length, URL, Optional, AnyOf, NumberRange

from internal.entity.dataset_entity import RetrievalStrategy
from internal.lib.helper import datetime_to_timestamp
from internal.model import Dataset, Document, DatasetQuery
from pkg.paginator.paginator import PaginatorReq


class CreateDataReq(FlaskForm):
    """创建知识库 http 请求校验"""
    name = StringField("name", validators=[
        DataRequired(message="知识库的名称不能为空"),
        Length(max=256, message="知识库名称长度不能超过256个字！")
    ])

    icon = StringField("icon", validators=[
        DataRequired(message="知识库的图标不能为空"),
        URL(message="知识库的图标必须是URL地址")
    ])

    description = StringField("description", default="",validators=[
        Optional(),
        Length(max=2000, message="知识库描述长度不能超过2000个字")
    ])


class GetDatasetResp(Schema):
    """获取知识库详情返回http响应结构"""
    id = fields.UUID(dump_default="")
    name = fields.String(dump_default="")
    icon = fields.String(dump_default="")
    description = fields.String(dump_default="")
    document_count = fields.Integer(dump_default=0)
    hit_count = fields.Integer(dump_default=0)
    related_app_count = fields.Integer(dump_default=0)
    character_count = fields.Integer(dump_default=0)
    updated_at = fields.Integer(dump_default=0)
    created_at = fields.Integer(dump_default=0)

    @pre_dump
    def process_data(self, data:Dataset, **kwargs):
        return {
            "id":data.id,
            "name":data.name,
            "icon":data.icon,
            "description":data.description,
            "document_count":data.document_count,
            "hit_count":data.hit_count,
            "related_app_count":data.related_app_count,
            "character_count":data.character_count,
            "created_at":int(data.updated_at.timestamp()),
            "updated_at":int(data.updated_at.timestamp()),
        }


class UpdateDatasetReq(FlaskForm):
    """更新知识库请求"""
    name = StringField("name", validators=[
        DataRequired(message="知识库的名称不能为空"),
        Length(max=256, message="知识库名称长度不能超过256个字！")
    ])

    icon = StringField("icon", validators=[
        DataRequired(message="知识库的图标不能为空"),
        URL(message="知识库的图标必须是URL地址")
    ])

    description = StringField("description", default="", validators=[
        Optional(),
        Length(max=2000, message="知识库描述长度不能超过2000个字")
    ])


class GetDatasetWithPageReq(PaginatorReq):
    search_word = StringField("search_word", validators=[Optional()])


class GetDatasetWithPageResp(Schema):
    """获取知识库分页列表响应数据"""
    id = fields.UUID(dump_default="")
    name = fields.String(dump_default="")
    icon = fields.String(dump_default="")
    description = fields.String(dump_default="")
    document_count = fields.Integer(dump_default=0)
    hit_count = fields.Integer(dump_default=0)
    related_app_count = fields.Integer(dump_default=0)
    character_count = fields.Integer(dump_default=0)
    updated_at = fields.Integer(dump_default=0)
    created_at = fields.Integer(dump_default=0)

    @pre_dump
    def process_data(self, data: Dataset, **kwargs):
        return {
            "id": data.id,
            "name": data.name,
            "icon": data.icon,
            "description": data.description,
            "document_count": data.document_count,
            "hit_count": data.hit_count,
            "related_app_count": data.related_app_count,
            "character_count": data.character_count,
            "created_at": int(data.updated_at.timestamp()),
            "updated_at": int(data.updated_at.timestamp()),
        }


class UpdateDocumentNameReq(FlaskForm):
    """更新文档名称/基础信息请求"""
    name = StringField("name", validators=[
        DataRequired("文档名称不能为空"),
        Length(max=100, message="文档的名称长度不能超过100")
    ])


class GetDocumentsWithPageReq(PaginatorReq):
    """获取文档分页列表请求"""
    search_word = StringField("search_word", default="", validators=[
        Optional()
    ])


class GetDocumentResp(Schema):
    """获取文档基础信息响应结构"""
    id = fields.UUID(dump_default="")
    dataset_id = fields.UUID(dump_default="")
    name = fields.String(dump_default="")
    segment_count = fields.Integer(dump_default=0)
    character_count = fields.Integer(dump_default=0)
    hit_count = fields.Integer(dump_default=0)
    position = fields.Integer(dump_default=0)
    enabled = fields.Bool(dump_default=False)
    disabled_at = fields.Integer(dump_default=0)
    status = fields.String(dump_default="")
    error = fields.String(dump_default="")
    updated_at = fields.Integer(dump_default=0)
    created_at = fields.Integer(dump_default=0)

    @pre_dump
    def process_data(self, data: Document, **kwargs):
        return {
            "id": data.id,
            "dataset_id": data.dataset_id,
            "name": data.name,
            "segment_count": data.segment_count,
            "character_count": data.character_count,
            "hit_count": data.hit_count,
            "position": data.position,
            "enabled": data.enabled,
            "disabled_at": datetime_to_timestamp(data.disabled_at),
            "status": data.status,
            "error": data.error,
            "updated_at": datetime_to_timestamp(data.updated_at),
            "created_at": datetime_to_timestamp(data.created_at),
        }


class GetDocumentsWithPageResp(Schema):
    """获取文档分页列表响应结构"""
    id = fields.UUID(dump_default="")
    name = fields.String(dump_default="")
    character_count = fields.Integer(dump_default=0)
    hit_count = fields.Integer(dump_default=0)
    position = fields.Integer(dump_default=0)
    enabled = fields.Bool(dump_default=False)
    disabled_at = fields.Integer(dump_default=0)
    status = fields.String(dump_default="")
    error = fields.String(dump_default="")
    updated_at = fields.Integer(dump_default=0)
    created_at = fields.Integer(dump_default=0)

    @pre_dump
    def process_data(self, data: Document, **kwargs):
        return {
            "id": data.id,
            "name": data.name,
            "character_count": data.character_count,
            "hit_count": data.hit_count,
            "position": data.position,
            "enabled": data.enabled,
            "disabled_at": datetime_to_timestamp(data.disabled_at),
            "status": data.status,
            "error": data.error,
            "updated_at": datetime_to_timestamp(data.updated_at),
            "created_at": datetime_to_timestamp(data.created_at),
        }


class HitReq(FlaskForm):
    """知识库召回请求 参数 """
    query = StringField("query", validators=[
        DataRequired("查询语句不能为空！"),
        Length(max=1000, message="查询语句最大长度不能超过1000个字！")
    ])
    retrieval_strategy = StringField("retrieval_strategy", validators=[
        DataRequired("检索策略不能为空！"),
        AnyOf([item.value for item in RetrievalStrategy], message="检索格式错误")
    ])
    k = IntegerField("k", validators=[
        DataRequired(message="最大召回数量不能为空"),
        NumberRange(min=1, max=10, message="最大召回数1-10")
    ])
    score = FloatField("score", validators=[
        NumberRange(min=0, max=0.99, message="最小匹配范围因该是0-0.99")
    ])



class GetDatasetQueriesResp(Schema):
    """获取知识库最近查询响应结构"""
    id = fields.UUID(dump_default="")
    dataset_id = fields.UUID(dump_default="")
    query = fields.String(dump_default="")
    source = fields.String(dump_default="")
    created_at = fields.Integer(dump_default=0)
    @pre_dump
    def process_data(self, data: DatasetQuery, **kwargs):
        return {
            "id": data.id,
            "dataset_id": data.dataset_id,
            "query": data.query,
            "source": data.source,
            "created_at": datetime_to_timestamp(data.created_at),
        }
